"""
!/usr/bin/env python
-*- coding: utf-8 -*-
@CreateTime : 2024/7/2 15:38
@Author  :    AnimateX
@Contact :    animatex@163.com
@File    :    fast_precision_test.py
@License :    Copyright © 2024 AnimateX. All rights reserved.
@Version :    fast_precision_test_VER0.1

------------------------------------------------------------------
# @Description:

------------------------------------------------------------------
"""
import os
import cv2
from tqdm import tqdm
from natsort import natsorted
import matplotlib.pyplot as plt
from fast_precision_core import FastPrecision
from parser_cam_param import ParserCamParam
from utils import Utils


def preprocess(json_path: str, test_precision_dir: str, disp_w, disp_h, key='_depth_', en_disp=False):
    if not os.path.exists(json_path):
        print(' [ERROR] Check json file path.')
        return

    if not os.path.exists(test_precision_dir):
        print(' [ERROR] Test directory is empty')
        return
    parser = ParserCamParam(json_path)
    cam_param, ratio_w, ratio_h = parser.parser_cam_param(disp_w, disp_h)

    test_img_name_list = os.listdir(test_precision_dir)

    img_name_list = [img_name for img_name in test_img_name_list if key in img_name and '.png' in img_name]
    img_path_list = [os.path.join(test_precision_dir, img_name) for img_name in img_name_list]
    img_path_list = natsorted(img_path_list)
    if key == '_disp' or key == '_disp_':
        en_disp = True

    if en_disp:
        dp_img_path_list = []
        for idx, disp_img_path in enumerate(tqdm(img_path_list, desc="Disp to Depth...", colour='cyan')):
            dp_img_path = disp_img_path.replace(key, '_depth_')
            disp_img = cv2.imread(disp_img_path, cv2.IMREAD_UNCHANGED)
            dp_img = Utils().disp_to_depth(disp_img,
                                           cam_param['q_coeffs']['Q23'],
                                           cam_param['q_coeffs']['Q32'],
                                           cam_param['q_coeffs']['Q33'],
                                           sub_pixel_value=64,
                                           zoom_ratio=1.0,
                                           max_dis=3000)
            cv2.imwrite(dp_img_path, dp_img)
            dp_img_path_list.append(dp_img_path)

        return cam_param, dp_img_path_list, ratio_w, ratio_h
    else:
        return cam_param, img_path_list, ratio_w, ratio_h


def test_f_type_precision(depth_img_path_list: list, depth_intrinsic: dict, ratio_w, ratio_h, denoise=False):
    if len(depth_img_path_list) == 0:
        print(' [ERROR] List is empty!')
        return

    test_type = 'f-type'
    delta = 10

    fig, axs = plt.subplots(1, len(depth_img_path_list), figsize=(7 * len(depth_img_path_list), 5))

    depth_img_path_list = natsorted(depth_img_path_list)
    for idx, dp_img_path in enumerate(tqdm(depth_img_path_list, desc=' Processing F-Type...', colour='green')):
        depth_img_name = dp_img_path.split('/')[-1]

        key_str = depth_img_name.split('_')[-1].split('.')[0]
        base_distance = int(depth_img_name.split('.')[0].split('_')[-1].split('-')[0])
        cam_distance = int(depth_img_name.split('.')[0].split('_')[-1].split('-')[1].replace('mm', ''))

        if depth_intrinsic is None or len(depth_intrinsic) == 0:
            depth_intrinsic = {
                'fx': 622.810913,
                'fy': 622.810913,
                'cx': 634.000000,
                'cy': 397.000000,
                'width': 1280,
                'height': 800
            }

        roi_dict = {
            '200-400mm': {
                's_y': int(300 * ratio_h),
                'e_y': int(580 * ratio_h),
                's_x': int(240 * ratio_w),
                'e_x': int(1279 * ratio_w)
            },
            '400-600mm': {
                's_y': int(300 * ratio_h),
                'e_y': int(530 * ratio_h),
                's_x': int(240 * ratio_w),
                'e_x': int(1279 * ratio_w)
            },
            '600-800mm': {
                's_y': int(300 * ratio_h),
                'e_y': int(503 * ratio_h),
                's_x': int(470 * ratio_w),
                'e_x': int(1279 * ratio_w)
            },
            '800-1000mm': {
                's_y': int(300 * ratio_h),
                'e_y': int(475 * ratio_h),
                's_x': int(433 * ratio_w),
                'e_x': int(1279 * ratio_w)
            }
        }

        roi_index_dict = roi_dict[key_str]

        fast_precision = FastPrecision(depth_intrinsic)
        point_cloud, zero_cnt = fast_precision.typePrecision(dp_img_path,
                                                             cam_distance,
                                                             depth_intrinsic['width'],
                                                             depth_intrinsic['height'],
                                                             roi_index_dict,
                                                             base_dis=base_distance,
                                                             delta=delta,
                                                             denoise=denoise,
                                                             test_type=test_type)

        # [05] gen image
        if test_type == 'f-type':
            title_str = f'F-Type BaseDis: {base_distance:4d}mm, CameraDis: {cam_distance:4d}mm'
            Utils().plot_point_cloud_side_view(point_cloud, base_distance, cam_distance, delta, title_str, ax=axs[idx],
                                               en_single_line=False)
        elif test_type == 'l-type':
            title_str = f'L-Type CameraDis: {cam_distance:4d}mm'
            Utils().plot_point_cloud_side_view(point_cloud, base_distance, cam_distance, delta, title_str, ax=axs[idx],
                                               en_single_line=True)
        elif test_type == 'u-type':
            title_str = f'U-Type CameraDis: {cam_distance:4d}mm'
            Utils().plot_point_cloud_side_view(point_cloud, base_distance, cam_distance, delta, title_str, ax=axs[idx],
                                               en_single_line=True)
    plt.tight_layout()
    # plt.savefig('F-TypeTestRes.png')
    plt.show()


def test_l_type_precision(depth_img_path_list: list, depth_intrinsic: dict, ratio_w, ratio_h, denoise=False):
    if len(depth_img_path_list) == 0:
        print(' [ERROR] List is empty!')
        return

    test_type = 'l-type'
    delta = 10

    fig, axs = plt.subplots(1, len(depth_img_path_list), figsize=(7 * len(depth_img_path_list), 5))

    depth_img_path_list = natsorted(depth_img_path_list)
    for idx, dp_img_path in enumerate(tqdm(depth_img_path_list, desc=' Processing L-Type...', colour='green')):
        depth_img_name = dp_img_path.split('/')[-1]

        key_str = depth_img_name.split('_')[-1].split('.')[0]
        base_distance = 0
        cam_distance = int(key_str.replace('mm', ''))

        if depth_intrinsic is None or len(depth_intrinsic) == 0:
            depth_intrinsic = {
                'fx': 622.810913,
                'fy': 622.810913,
                'cx': 634.000000,
                'cy': 397.000000,
                'width': 1280,
                'height': 800
            }

        roi_dict = {
            '200mm': {
                's_y': int(180 * ratio_h),
                'e_y': int(590 * ratio_h),
                's_x': int(0 * ratio_w),
                'e_x': int(1279 * ratio_w)
            },
            '400mm': {
                's_y': int(180 * ratio_h),
                'e_y': int(590 * ratio_h),
                's_x': int(0 * ratio_w),
                'e_x': int(1279 * ratio_w)
            },
            '600mm': {
                's_y': int(270 * ratio_h),
                'e_y': int(530 * ratio_h),
                's_x': int(0 * ratio_w),
                'e_x': int(1279 * ratio_w)
            },
            '800mm': {
                's_y': int(300 * ratio_h),
                'e_y': int(500 * ratio_h),
                's_x': int(0 * ratio_w),
                'e_x': int(1279 * ratio_w)
            },
            '1000mm': {
                's_y': int(300 * ratio_h),
                'e_y': int(490 * ratio_h),
                's_x': int(0 * ratio_w),
                'e_x': int(1279 * ratio_w)
            }
        }

        roi_index_dict = roi_dict[key_str]

        fast_precision = FastPrecision(depth_intrinsic)
        point_cloud, zero_cnt = fast_precision.typePrecision(dp_img_path,
                                                             cam_distance,
                                                             depth_intrinsic['width'],
                                                             depth_intrinsic['height'],
                                                             roi_index_dict,
                                                             base_dis=base_distance,
                                                             delta=delta,
                                                             denoise=denoise,
                                                             test_type=test_type)

        if test_type == 'f-type':
            title_str = f'F-Type BaseDis: {base_distance:4d}mm, CameraDis: {cam_distance:4d}mm'
            Utils().plot_point_cloud_side_view(point_cloud, base_distance, cam_distance, delta, title_str, ax=axs[idx],
                                               en_single_line=False)
        elif test_type == 'l-type':
            title_str = f'L-Type CameraDis: {cam_distance:4d}mm'
            Utils().plot_point_cloud_side_view(point_cloud, base_distance, cam_distance, delta, title_str, ax=axs[idx],
                                               en_single_line=True)
        elif test_type == 'u-type':
            title_str = f'U-Type CameraDis: {cam_distance:4d}mm'
            Utils().plot_point_cloud_side_view(point_cloud, base_distance, cam_distance, delta, title_str, ax=axs[idx],
                                               en_single_line=True)
    plt.tight_layout()
    # plt.savefig('L-TypeTestRes.png')
    plt.show()


def test_u_type_precision(depth_img_path_list: list, depth_intrinsic: dict, ratio_w, ratio_h, denoise=False):
    if len(depth_img_path_list) == 0:
        print(' [ERROR] List is empty!')
        return

    test_type = 'u-type'
    delta = 10

    fig, axs = plt.subplots(1, len(depth_img_path_list), figsize=(7 * len(depth_img_path_list), 5))

    depth_img_path_list = natsorted(depth_img_path_list)
    for idx, dp_img_path in enumerate(tqdm(depth_img_path_list, desc=' Processing U-Type...', colour='green')):
        depth_img_name = dp_img_path.split('/')[-1]

        key_str = depth_img_name.split('_')[-1].split('.')[0]
        base_distance = 0
        cam_distance = int(key_str.replace('mm', ''))

        if depth_intrinsic is None or len(depth_intrinsic) == 0:
            depth_intrinsic = {
                'fx': 622.810913,
                'fy': 622.810913,
                'cx': 634.000000,
                'cy': 397.000000,
                'width': 1280,
                'height': 800
            }

        roi_dict = {
            '200mm': {
                's_y': int(180 * ratio_h),
                'e_y': int(590 * ratio_h),
                's_x': int(0 * ratio_w),
                'e_x': int(1279 * ratio_w)
            },
            '400mm': {
                's_y': int(180 * ratio_h),
                'e_y': int(590 * ratio_h),
                's_x': int(0 * ratio_w),
                'e_x': int(1279 * ratio_w)
            },
            '600mm': {
                's_y': int(270 * ratio_h),
                'e_y': int(530 * ratio_h),
                's_x': int(0 * ratio_w),
                'e_x': int(1279 * ratio_w)
            },
            '800mm': {
                's_y': int(300 * ratio_h),
                'e_y': int(500 * ratio_h),
                's_x': int(100 * ratio_w),
                'e_x': int(1279 * ratio_w)
            },
            '1000mm': {
                's_y': int(300 * ratio_h),
                'e_y': int(490 * ratio_h),
                's_x': int(0 * ratio_w),
                'e_x': int(1279 * ratio_w)
            }
        }

        roi_index_dict = roi_dict[key_str]

        fast_precision = FastPrecision(depth_intrinsic)
        point_cloud, zero_cnt = fast_precision.typePrecision(dp_img_path,
                                     cam_distance,
                                     depth_intrinsic['width'],
                                     depth_intrinsic['height'],
                                     roi_index_dict,
                                     base_dis=base_distance,
                                     delta=delta,
                                     denoise=denoise,
                                     test_type=test_type)

        if test_type == 'f-type':
            title_str = f'F-Type BaseDis: {base_distance:4d}mm, CameraDis: {cam_distance:4d}mm'
            Utils().plot_point_cloud_side_view(point_cloud, base_distance, cam_distance, delta, title_str, ax=axs[idx],
                                               en_single_line=False)
        elif test_type == 'l-type':
            title_str = f'L-Type CameraDis: {cam_distance:4d}mm'
            Utils().plot_point_cloud_side_view(point_cloud, base_distance, cam_distance, delta, title_str, ax=axs[idx],
                                               en_single_line=True)
        elif test_type == 'u-type':
            title_str = f'U-Type CameraDis: {cam_distance:4d}mm'
            Utils().plot_point_cloud_side_view(point_cloud, base_distance, cam_distance, delta, title_str, ax=axs[idx],
                                               en_single_line=True)

    plt.tight_layout()
    # plt.savefig('U-TypeTestRes.png')
    plt.show()


def test_fast_precision(fast_precision_path: str, json_path: str, disp_image_w, disp_image_h, test_type='f-type',
                        key='_disp_', en_disp=False):
    test_dir_path = os.path.join(fast_precision_path, test_type)
    cam_param, dp_img_path_list, ratio_w, ratio_h = preprocess(json_path, test_dir_path, disp_image_w, disp_image_h,
                                                               key=key, en_disp=en_disp)

    if test_type == 'f-type':
        test_f_type_precision(dp_img_path_list, cam_param['depth_intrinsics'], ratio_w, ratio_h, denoise=True)
    elif test_type == 'l-type':
        test_l_type_precision(dp_img_path_list, cam_param['depth_intrinsics'], ratio_w, ratio_h, denoise=True)
    elif test_type == 'u-type':
        test_u_type_precision(dp_img_path_list, cam_param['depth_intrinsics'], ratio_w, ratio_h, denoise=True)


if __name__ == '__main__':
    fast_precision_dir_path = ("/home/dataset_b/rmslTestDataset/orbbec/gemini335-01/v0619/"
                               "r1920x1080_d1280x800/rk_result_0925/precision/fast-precision/")
    json_file_path = "/home/dataset_b/rmslTestDataset/orbbec/gemini335-01/v0619/rgb1920x1080_ir1280x800.json"

    disp_img_w = 640
    disp_img_h = 400
    mode = '_disp_'
    en_disp = True

    # disp_img_w = 640
    # disp_img_h = 480

    """
        [Note]: Need to alter!
            [01] 快速精度数据集地址
            [02] 模组对应测试集分辨率对应的内外参以及旋转矩阵和平移系数对应的参数json
            [03] 测试的快速精度子类别
            [04] rk生成的深度需要修改关键字为 '_xxx_depth_' 后缀保持和默认一致
    """
    test_fast_precision(fast_precision_dir_path, json_file_path, disp_img_w, disp_img_h, test_type='f-type',
                        key=mode, en_disp=en_disp)
    test_fast_precision(fast_precision_dir_path, json_file_path, disp_img_w, disp_img_h, test_type='l-type',
                        key=mode, en_disp=en_disp)
    test_fast_precision(fast_precision_dir_path, json_file_path, disp_img_w, disp_img_h, test_type='u-type',
                        key=mode, en_disp=en_disp)
